%-------------------------------------------------------------------------%    
%This program solves the model with the TSP plan for 
%workers with frozen DB plans. Freeze age =60.
% 1% unconditional match
% 1-1 match per dollar contributed upto 3%
% 0.5-1 match per dollar contributed for the next 2%
% 5% cap
%-------------------------------------------------------------------------% 

cd '/Users/dpatki/Dropbox/Research/Aging LFP/structural model/programs';
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/Disclosed output/HRS MiCDA/20210405/Tables');
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/Disclosed output/20190715/graphing');
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/data');
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/structural model/programs/lininterp'); 
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/structural model/programs/v-funcs-testing'); 

%Parameters from estimated model
param = csvread('/Users/dpatki/Dropbox/Research/Aging LFP/output/x1_sobolstrt.csv');

        %=========================================================================%
        % Part 1: parameters
        %=========================================================================%
        rng(1,'twister');
        s = rng;
 	    sigma = param(1);      %EIS      
	    beta  = .97;           %discount factor 
        aT    = 81;            %terminal age
        aS    = 50;            %start age
        r     = .029;          %interest rate 
        pi    = .028;          %inflation
        V_term = 0;            %terminal value
        z_age = 60;            %freeze age
        hor = aT-1-z_age;      %horizon (years in model post freeze)
        N_obs  = 5000;         %Number of simulated individuals
        ages = aS:1:aT-1;      %ages
        t = (1:(aT-aS))';      %time periods
        scaler = 1;            %scaler on consumption utility 
        solver_param = [sigma;beta;aT;aS;r;V_term;scaler;pi]; %collect some params
      
        %-------------------------------------------------------------------------%
        %Calculate annuity value and earnings per period from HRS
        %-------------------------------------------------------------------------%
        %1. HRS db wealth by age
        pdata = csvread('compwealth_frz_w10_22.csv',1,0);
        %col2 is age
        %col3 is avg db wealth in 2010 dollars
        %col4 is avg earnings in 2010 dollars
        dbw = pdata(aS-48+1:aT-48,2:3);
        y_raw = pdata(aS-48+1:end,4);
        %Smooth the earnings profile using cubic polynomic fit
        y_fit_param =  polyfit(ages',y_raw,3);
        y = polyval(y_fit_param,ages');
        %percent earnings losses from freezes (periods 0-4)
        y_cut = [.037;.020;.015;.048;.016]; %[0; 0; 0; 0; 0]; %
        %2. SS annuity by retirement age (hh-level)
        b_ss = csvread('ssw10.csv',1,0);
        b_ss = b_ss(aS-48+1:end,4);
        %3. SSA mortality rates by age (48+)
        pdeath = csvread('SSA_mortality_2010.csv',1,0);
        pdeath = pdeath(aS+1:end,:);
        avpdeath = mean(pdeath(:,2:3),2);
        %4. annuity factors
        afactor = zeros(length(avpdeath),1); 
        %loop over age
        for a = 1:length(afactor)
            disc = zeros(length(afactor)-a+1,1);
            disc(1) = 1;
            for i = 2:length(disc)
                disc(i) = (1-avpdeath(a+i-1))*(1/((1+r)*(1+pi)))^(i-1); 
            end

            afactor(a) = sum(disc);
        end     
        afactor = [pdeath(1:aT-aS,1) afactor(1:aT-aS)];
        %DB annuity by retirement age
        b_db = dbw(:,2)./afactor(:,2);
        %total annuity benefit
        b = b_db + b_ss;
       
        %-------------------------------------------------------------------------%
        % Normalized mortality rate
        %-------------------------------------------------------------------------%
        mortp = avpdeath(1:aT-aS);
        mortp_80 = mortp./avpdeath(aT-aS+1);
        
        %-------------------------------------------------------------------------%
        % Asset grids (A = non-pension; W = DC pension)
        %-------------------------------------------------------------------------%
        %A
        A_gridN = 20;
        split = 5;
        A_mins = [0;50000;150000;500000;1000000;2000000];
        As=zeros(A_gridN/split,split);
        for i = 1:5
            step = (A_mins(i+1)-A_mins(i))/(A_gridN/split);
            As(:,i) = A_mins(i):step:(A_mins(i+1)-step);
        end
        A = reshape(As,[A_gridN,1]);
        
        %W
        W_gridN = 20;
        split = 5;
        W_mins = [0;25000;75000;250000;500000;1500000];
        Ws=zeros(W_gridN/split,split);
        for i = 1:5
            step = (W_mins(i+1)-W_mins(i))/(W_gridN/split);
            Ws(:,i) = W_mins(i):step:(W_mins(i+1)-step);
        end
        W = reshape(Ws,[W_gridN,1]);    
        
        %-------------------------------------------------------------------------%
        % Average tax rates from nber taxsim (age-dependent grids)
        %-------------------------------------------------------------------------%       
        worktax = csvread('workinc_taxsim_a_20.csv',0,0);
        %col1 = age, col2 = earn_before_dc, col3 =fica, col4=earn_after_dc
        %col5=fiitax, col6=dc_contrib
        itau_w = worktax(:,5)./worktax(:,4); %income tax
        ptau_w = 0.5*worktax(:,3)./worktax(:,2); %worker share of FICA
      
        rettax = csvread('retinc_taxsim_a_20.csv',0,0);
        %col1 = fiitax, col2 = age, col3=pension distribution, col4 = SS
        itau_r = rettax(:,1)./(rettax(:,3)+rettax(:,4));
             
        %Collect parameters for model solution
        taxf = cell(3,1);
        taxf{1,1}=itau_w;
        taxf{2,1}=ptau_w;
        taxf{3,1}=itau_r;

        %freeze-affected average tax rates 
        taxfz = cell(9,3);
        zindex = 1;
        for aa = 56:64         
            i_z = find(ages==aa);
            
            %While working
            fname = strcat('workinc_taxsim_a_20_tspz',num2str(i_z),'.csv');
            worktax_z = csvread(fname,0,0);
            itau_w_z = worktax_z(:,5)./worktax_z(:,4); %income tax
            ptau_w_z = 0.5*worktax_z(:,3)./worktax_z(:,2); %worker share of FICA
            taxfz{zindex,1} = itau_w_z;
            taxfz{zindex,2} = ptau_w_z;
     
            %Retirement
            fname = strcat('retinc_taxsim_a_20_z',num2str(i_z),'.csv');
            rettax_z = csvread(fname,0,0);
            itau_r_z = rettax_z(:,1)./(rettax_z(:,3)+rettax_z(:,4));
            taxfz{zindex,3} = itau_r_z;
            zindex = zindex+1;
        end

        %-------------------------------------------------------------------------%
        %Disutility of work process 
        %-------------------------------------------------------------------------%
    	mu      = 0;  %mean normal underlying f
        rho     = 1/(1+exp(param(2)));  %persistence of random process 
    	sigma_v = exp(param(3)); %sd of normal underlying f
    	phi     = param(4); %age slope
        gamma   = param(5); %constant
    	f_N     = 6; %grid size

        %=========================================================================%
        % Part 2: solve model for value and policy function 
        % without any freeze activity
        %=========================================================================%

        %Cell array to store value and policy functions for each age
        optfunc = cell(length(t),8);

        %-------------------------------------------------------------------------%
        %Solve age T-1 policy and value functions outside the loop
        %-------------------------------------------------------------------------%	
        %Set last period=true
        last=1;
        a=0;
        [V_w,A_polw,W_polw]=vwork_5(A,W,last,a,y,optfunc,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxf,mortp_80);
        [V_r,A_polr,W_polr]=vret_4(A,W,last,a,b_db,b_ss,optfunc,solver_param,taxf,mortp_80);
	
        
        % Outer value function and labor supply choice
        V = zeros(size(V_w)); 
        L = zeros(size(V_w)); 
       
        for i = 1:length(A)
            for j = 1:length(W)
                for k = 1:f_N
                    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
                    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
                end
            end
        end

        %Store value and policy functions
        optfunc(1,:) = {V_w;A_polw;W_polw;V_r;A_polr;W_polr;V;L};
        
        %-------------------------------------------------------------------------%
        %Solve policy and value functions for remaining periods
        %-------------------------------------------------------------------------%      
        %Set last period = false
        last=0;
        
        %Loop through ages
        for a = 1:length(t)-1 
    

               [V_w,A_polw,W_polw]=vwork_5(A,W,last,a,y,optfunc,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxf,mortp_80);
               [V_r,A_polr,W_polr]=vret_4(A,W,last,a,b_db,b_ss,optfunc,solver_param,taxf,mortp_80);


		% Outer value function and labor supply choice
		V = zeros(size(V_w)); 
		L = zeros(size(V_w)); 

		for i = 1:length(A)
		    for j = 1:length(W)
			for k = 1:f_N
			    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
			    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
			end
		    end
		end

		%Store value and policy functions
        optfunc(a+1,:) = {V_w;A_polw;W_polw;V_r;A_polr;W_polr;V;L};

        end
        

        %=========================================================================%
        % Part 3: solve model for value and policy function 
        % separately with counterfactual DC/pay policies
        %=========================================================================%	
        %Cell array to store value and policy functions for each freeze age
        % dimension is age x freeze ages x number of functions (8)
        optfunc_frz = cell(length(t),8);
        lastfuncW0 = [];
        lastfuncR0 = [];
       
           
        %Freeze the path of DB pension annuities (keep SS wealth fixed)
        z = 5;
        i_z = find(ages==z_age-1);
        b_db_z = b_db;
        b_db_z(i_z:end) = b_db(i_z);
      
        %freeze fixes nominal annuity; compute real annuity at each
        %prospective retirement age
        for i = i_z:length(b_db_z) 
            b_db_z(i) = b_db_z(i)*(1/(1+pi))^(i-i_z);
        end
        
        b_z = b_db_z + b_ss;
        %earn profile given freeze at age z
        y_loss = ones(length(y),1);
        y_loss(i_z+1:i_z+1+length(y_cut)-1) = 1-y_cut;
        y_z = y.*y_loss;

        %---------------------------------------------------------------------%
        %Solve age T-1 policy and value functions outside the age loop
        %---------------------------------------------------------------------%
        %Set last period=true
        last=1;
        a=0;

        [V_w,A_polw,W_polw]=vworkz_tsp(A,W,last,a,y_z,lastfuncW0,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxfz,mortp_80,z);
        [V_r,A_polr,W_polr]=vretz_4(A,W,last,a,b_db_z,b_ss,lastfuncR0,solver_param,taxfz,mortp_80,z);


        % Outer value function and labor supply choice
        V = zeros(size(V_w)); 
        L = zeros(size(V_w)); 

        for i = 1:length(A)
            for j = 1:length(W)
                for k = 1:f_N
                    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
                    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
                end
            end
        end

        %Store value and policy functions
        optfunc_frz(1,:) = {V_w,A_polw,W_polw,V_r,A_polr,W_polr,V,L};   

        %-------------------------------------------------------------------------%
        %Solve policy and value functions for remaning periods
        %-------------------------------------------------------------------------%
        %Set last period = false
        last=0;

        %Loop through ages
        for a = 1:length(t)-1 

            lastfuncW = cell2mat(optfunc_frz(a,7));
            lastfuncR = optfunc_frz(a,4);

            [V_w,A_polw,W_polw]=vworkz_tsp(A,W,last,a,y_z,lastfuncW,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxfz,mortp_80,z);
            [V_r,A_polr,W_polr]=vretz_4(A,W,last,a,b_db_z,b_ss,lastfuncR,solver_param,taxfz,mortp_80,z);


            % Outer value function and labor supply choice
            V = zeros(size(V_w)); 
            L = zeros(size(V_w)); 

            for i = 1:length(A)
                for j = 1:length(W)
                for k = 1:f_N
                    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
                    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
                end
                end
            end

            %Store value and policy functions
            optfunc_frz(a+1,:) = {V_w;A_polw;W_polw;V_r;A_polr;W_polr;V;L};

        end
      
        %=========================================================================%
        %Part 4: Simulate shocks and solve for labor supply and saving
        %=========================================================================%
        
        %Cell arrays to hold simulated paths 
        LsimC = cell(1,1);
        AsimC = cell(1,1);
        WsimC = cell(1,1);
        LsimT = cell(1,1);
        AsimT = cell(1,1);
        WsimT = cell(1,1);
      
        %Nodes of discrete AR(1) process -- need max and min of these to
        %assign values out of bounds in the simulation
        [f,~] = rouwenhorst(f_N,mu,rho,sigma_v);
        
        %Simulated observations 
        obs = N_obs;
        
        %Disutility shock process
        %simulate AR(1) shocks. 1000 period burn-in... 
        f_sim0 = zeros(1000+(hor+6),N_obs);
        rng(s);
        v = normrnd(0,sigma_v,[length(f_sim0),N_obs]);
        for i = 2:length(f_sim0)
            f_sim0(i,:) = (1-rho)*mu+rho*f_sim0(i-1,:)+v(i,:);
        end
        %Trim away burn-in periods
        f_sim0 = f_sim0(end-(hor+6-1):end,:);
        D{1,1} = f_sim0;
       
        %Initial assets five years prior to frz (age 55) --moments from HRS
        a=5;
        A0 = cell(1,1);
        W0 = cell(1,1);
        A_params = csvread('A_params.csv',1,0);
        J_params = csvread('J_params.csv',1,0);
        rng(s);    
        %Number of individuals with 0 DC wealth
        frac0obs = round(A_params(a,4)*N_obs); 
        %Non-DC wealth for individuals with 0 DC wealth
        A_temp0 = lognrnd(A_params(a,2),A_params(a,3),[frac0obs,1]);
        %DC wealth for individuals with no DC balances = 0
        W_temp0 = zeros(frac0obs,1);
        %Joint wealth distribution for individuals with >0 DC wealth
        mvnmu = J_params(a,2:3);
        mvnsig = [(J_params(a,4))^2 J_params(a,6);J_params(a,6) (J_params(a,5))^2];
        J_wealth = exp(mvnrnd(mvnmu,mvnsig,(N_obs-frac0obs)));
        W_temp1 = J_wealth(:,1);
        A_temp1 = J_wealth(:,2);
        A0{1,1} = [A_temp0' A_temp1'];
        W0{1,1} = [W_temp0' W_temp1'];

        %use policy functions to obtain optimal labor supply and asset
        %accumulation for each simulated individual in each period 
        %-------------------------------------------------------------------------%
        %Control group
        %-------------------------------------------------------------------------%
        
        p=5; %Age 55 in period -5
        
        L_sim = zeros(hor+6,N_obs);
        A_sim = zeros(hor+6,N_obs);
        W_sim = zeros(hor+6,N_obs);
        A_sim(1,:) = A0{1,1};
        W_sim(1,:) = W0{1,1};

        %Set disutility distribution
        f_sim = D{1,1};
 
        %Employment choice for t = -5 outside the loop (t = period
        %relative to freeze) a = 1 when t = -5
        for i = 1:N_obs
            a=1;    
            Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
            Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
            Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
            Afunc_r  = Afunc_r1{p+a};
            Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
            Wfunc_r  = Wfunc_r1{p+a};
            Lfunc = optfunc{(aT-aS)-(p+a-1),8};
            %Put extreme draws of A,W,f onto the grid 
            %A
            if A_sim(a,i)>max(A)
                A_sim(a,i) = max(A);
            end
            %W
            if W_sim(a,i)>max(W)
                W_sim(a,i) = max(W);
            end	
            %f                
            if f_sim(a,i)<min(f) 
                f_sim(a,i) = min(f);
            elseif f_sim(a,i)>max(f)
                f_sim(a,i) = max(f);
            end      

            %Simulated values 
            L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));
            
            if L_sim(a,i) == 1
                A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
            elseif L_sim(a,i) == 0 
                A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
            end 
        end

        for a = 2:hor+5 
           for i = 1:N_obs

                Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end                    
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                end

            end
        end

        %Final period employment choice (A,W already determined)
        a = hor+6;
        for i = 1:N_obs

                Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                                     
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                end


        end

        LsimC{1,1} = L_sim;
        AsimC{1,1} = A_sim;
        WsimC{1,1} = W_sim;
    
    %-------------------------------------------------------------------------%
	%Treated (i.e. DB frozen) group starting at age 55
	%-------------------------------------------------------------------------%
    
        p=5;  %Age 55 in period -5
        L_sim = zeros(hor+6,N_obs);
        A_sim = zeros(hor+6,N_obs);
        W_sim = zeros(hor+6,N_obs);
        A_sim(1,:) = A0{1,1};
        W_sim(1,:) = W0{1,1};

        %Set disutility distribution
        f_sim = D{1,1};

        %Employment choice for t = -5 outside the loop (t = period
        %relative to sunset) a = 1 when t = -5
        for i = 1:N_obs
            a=1;    
            Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
            Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
            Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
            Afunc_r  = Afunc_r1{p+a};
            Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
            Wfunc_r  = Wfunc_r1{p+a};
            Lfunc = optfunc{(aT-aS)-(p+a-1),8};
            %Put extreme draws of A,W,f onto the grid 
            %A
            if A_sim(a,i)>max(A)
                A_sim(a,i) = max(A);
            end
            %W
            if W_sim(a,i)>max(W)
                W_sim(a,i) = max(W);
            end	
            %f                
            if f_sim(a,i)<min(f) 
                f_sim(a,i) = min(f);
            elseif f_sim(a,i)>max(f)
                f_sim(a,i) = max(f);
            end      

            %Simulated values 
            L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));
            
            if L_sim(a,i) == 1
                A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
            elseif L_sim(a,i) == 0 
                A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
            end  
        end
        
       for a = 2:5 
           for i = 1:N_obs

                Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end                    
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                end

            end
       end

       for a = 6:hor+5 %after freeze
           for i = 1:N_obs

                Afunc_w = optfunc_frz{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc_frz{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc_frz{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc_frz{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc_frz{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end                    
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                end

            end
       end
       
        %Final period employment choice (A,W already determined)
        a = hor+6;
        for i = 1:N_obs

                Afunc_w = optfunc_frz{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc_frz{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc_frz{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc_frz{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc_frz{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                                     
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                end

        end

        LsimT{1,1} = L_sim;
        AsimT{1,1} = A_sim;
        WsimT{1,1} = W_sim;

    %-------------------------------------------------------------------------%
    % Compute moments by averaging over simulated individuals
    %-------------------------------------------------------------------------%
     
    %Concatenate control group       
    L_C2 = LsimC{1,1};
    A_C2 = AsimC{1,1};
    W_C2 = WsimC{1,1};
       
    %Condition on sample of workers who are still in LF 1 years pre-freeze
    indx_1 = find(L_C2(1,:)==0);

    L_C2_r1 = L_C2;
    L_C2_r1(:,indx_1)=[];
    A_C2_r1 = A_C2;
    A_C2_r1(:,indx_1)=[];
    W_C2_r1 = W_C2;
    W_C2_r1(:,indx_1)=[];

    %Extract treated group data               
    L_T2 = LsimT{1,1};
    A_T2 = AsimT{1,1};
    W_T2 = WsimT{1,1};    
  
    %Condition on sample of workers who are still in LF 1 years
    %pre-comparison
    indx_2 = find(L_T2(1,:)==0);

    L_T2_r1 = L_T2;
    L_T2_r1(:,indx_2)=[];
    A_T2_r1 = A_T2;
    A_T2_r1(:,indx_2)=[];
    W_T2_r1 = W_T2;
    W_T2_r1(:,indx_2)=[]; 
    

    %-------------------------------------------------------------------------%    
    %Export data for analysis
    %-------------------------------------------------------------------------% 

    %Create some inputs
    aT    = 81;            %terminal age
    aS    = 50;            %start age
    pdata = csvread('compwealth_frz_w10_22.csv',1,0);
    ages = aS:1:aT-1;      %ages
    y_raw = pdata(aS-48+1:end,4); %raw earnings profile
    y_fit_param =  polyfit(ages',y_raw,3);
    y = polyval(y_fit_param,ages');  %smoothed earnings profile
    ernstart = 6; %start at age 55
    %SSA mortality rates by age (48+)
    pdeath = csvread('SSA_mortality_2010.csv',1,0);
    pdeath = pdeath(aS+1:end,:);
    avpdeath = mean(pdeath(:,2:3),2);
    mortp = avpdeath(1:aT-aS);
    mortp_80 = mortp./avpdeath(aT-aS+1);

    %------------------------------
    %Unfrozen group
    %------------------------------


    %Extract simulated data 
    L = L_C2;
    A = A_C2;
    W = W_C2;
    ages = (55:1:80)';
    g = D{1,1}+gamma+phi.*ages;

    %Remove non-workers in 1st period
    indx_1 = find(L(1,:)==0);
    L(:,indx_1)=[];
    A(:,indx_1)=[];
    W(:,indx_1)=[];
    g(:,indx_1)=[];

    %Create id's
    ids = zeros(size(L));
    for i = 1:length(L)
        ids(:,i) = i;
    end

    ids_vec = reshape(ids,[(aT-(z_age-5))*length(L),1]);
    LT = reshape(L,[(aT-(z_age-5))*length(L),1]);
    AT = reshape(A,[(aT-(z_age-5))*length(L),1]);
    WT = reshape(W,[(aT-(z_age-5))*length(L),1]);
    gT = reshape(g,[(aT-(z_age-5))*length(L),1]);

    %Earnings 
    ern = y(ernstart:end);
    ernrep = repmat(ern,[length(L),1]);
    %SS benefits
    ssben = b_ss(ernstart:end);
    ssbenrep = repmat(ssben,[length(L),1]);
    %DB benefits
    B_db = b_db(ernstart:end);
    B_dbrep = repmat(B_db,[length(L),1]);

    %Probability of death
    mort = mortp_80(ernstart:end);
    mortrep = repmat(mort,[length(L),1]);

    Cdata = [ids_vec gT LT AT WT ernrep B_dbrep ssbenrep mortrep];

    %------------------------------
    %Frozen group
    %------------------------------

    %Extract simulated data 
    L = L_T2;
    A = A_T2;
    W = W_T2;
    ages = (55:1:80)';
    g = D{1,1}+gamma+phi.*ages;

    %Remove non-workers in 1st period
    indx_1 = find(L(1,:)==0);
    L(:,indx_1)=[];
    A(:,indx_1)=[];
    W(:,indx_1)=[];
    g(:,indx_1)=[];

    %Create id's
    ids = zeros(size(L));
    for i = 1:length(L)
        ids(:,i) = i;
    end 

    ids_vec = reshape(ids,[(aT-(z_age-5))*length(L),1]);
    LT = reshape(L,[(aT-(z_age-5))*length(L),1]);
    AT = reshape(A,[(aT-(z_age-5))*length(L),1]);
    WT = reshape(W,[(aT-(z_age-5))*length(L),1]);
    gT = reshape(g,[(aT-(z_age-5))*length(L),1]);

    %Earnings 
    ern = y_z(ernstart:end);
    ernrep = repmat(ern,[length(L),1]);
    %SS benefits
    ssben = b_ss(ernstart:end);
    ssbenrep = repmat(ssben,[length(L),1]);
    %DB benefits
    B_db = b_db_z(ernstart:end);
    B_dbrep = repmat(B_db,[length(L),1]);

    %Probability of death
    mort = mortp_80(ernstart:end);
    mortrep = repmat(mort,[length(L),1]);

    Tdata = [ids_vec gT LT AT WT ernrep B_dbrep ssbenrep mortrep];

    %Stack control and treated data
    trt_flag = ones(length(Tdata),1);
    Cdata = [Cdata (1-trt_flag)];
    Tdata = [Tdata trt_flag];
    data  = [Cdata;Tdata];

    csvwrite('/Users/dpatki/Dropbox/Research/Aging LFP/data/simulated_profiles_z60_tsp.csv',data);




